import torch.nn as nn

class ChemNet(nn.Module):
	def __init__(self, in_sz, hd_sz):
		super(ChemNet, self).__init__()

		base = hd_sz
		self.encoder = nn.Sequential(
			nn.Linear(in_sz, base * 2),
			nn.ReLU(),
			nn.BatchNorm1d(base * 2),
			nn.Linear(base * 2, base),
			nn.ReLU(),
			)
		self.decoder1 = nn.Sequential(
			nn.BatchNorm1d(base),
			nn.Linear(base, base * 2),
			nn.ReLU(),
			nn.BatchNorm1d(base * 2),
			nn.Linear(base * 2, in_sz),
			nn.ReLU()
			)
		self.decoder2 = nn.Sequential(
			nn.BatchNorm1d(base),
			nn.Linear(base, base * 2),
			nn.ReLU(),
			nn.BatchNorm1d(base * 2),
			nn.Linear(base * 2, in_sz),
			nn.ReLU()
			)

	def forward(self, x, channel):
		z = self.encoder(x)
		if channel == 1:
			return self.decoder1(z), z
		else:
			return self.decoder2(z), z
			